#include "Matrix.h"
#include <cassert>
#include <iostream>
#include <Eigen/Dense>

using namespace SG::Algebra;

namespace SG::Algebra
{
    /** @brief 通过伴随矩阵法计算 2*2 矩阵的逆 */
    Matrix inverse2X2 (const Matrix& m)
    {
        Real_t j = m (0, 0) * m (1, 1) - m (0, 1) * m (1, 0);
        Matrix inv (2, 2);
        inv (0, 0) = m (1, 1) / j;
        inv (1, 1) = m (0, 0) / j;
        inv (0, 1) = -m (0, 1) / j;
        inv (1, 0) = -m (1, 0) / j;
        return inv;
    }

    /** @brief 通过伴随矩阵法计算 3*3 矩阵的逆 */
    Matrix inverse3X3 (const Matrix& m, Real_t& det)
    {
        Matrix n (3, 3);  // m伴随矩阵
        n (0, 0) = m (1, 1) * m (2, 2) - m (1, 2) * m (2, 1);
        n (1, 0) = m (1, 2) * m (2, 0) - m (1, 0) * m (2, 2);
        n (2, 0) = m (1, 0) * m (2, 1) - m (1, 1) * m (2, 0);
        n (0, 1) = m (2, 1) * m (0, 2) - m (2, 2) * m (0, 1);
        n (1, 1) = m (2, 2) * m (0, 0) - m (2, 0) * m (0, 2);
        n (2, 1) = m (2, 0) * m (0, 1) - m (2, 1) * m (0, 0);
        n (0, 2) = m (0, 1) * m (1, 2) - m (0, 2) * m (1, 1);
        n (1, 2) = m (0, 2) * m (1, 0) - m (0, 0) * m (1, 2);
        n (2, 2) = m (0, 0) * m (1, 1) - m (0, 1) * m (1, 0);                   // 代数余子式

        det = m (0, 0) * n (0, 0) + m (0, 1) * n (1, 0) + m (0, 2) * n (2, 0);  // m行列式

        Matrix inv (3, 3);                                                      // m逆矩阵
        inv = n * (1 / det);
        return inv;
    }

    Real_t det2X2 (const Matrix& m)
    {
        // 先断定m的尺寸，否则报错
        if (m.GetRowSize () != 2 && m.GetColSize () != 2)
        {
            throw std::runtime_error ("Error: The matrix size is not 2*2");
        }
        return m (0, 0) * m (1, 1) - m (1, 0) * m (0, 1);
    }

    /** @brief 通过伴随矩阵法计算 3*3 矩阵的行列式 */
    Real_t det3X3 (const Matrix& m)
    {
        // 先断定m的尺寸，否则报错
        if (m.GetRowSize () != 3 && m.GetColSize () != 3)
        {
            throw std::runtime_error ("Error: The matrix size is not 3*3");
        }

        Matrix n (3, 3);  // m伴随矩阵
        n (0, 0) = m (1, 1) * m (2, 2) - m (1, 2) * m (2, 1);
        n (1, 0) = m (1, 2) * m (2, 0) - m (1, 0) * m (2, 2);
        n (2, 0) = m (1, 0) * m (2, 1) - m (1, 1) * m (2, 0);
        n (0, 1) = m (2, 1) * m (0, 2) - m (2, 2) * m (0, 1);
        n (1, 1) = m (2, 2) * m (0, 0) - m (2, 0) * m (0, 2);
        n (2, 1) = m (2, 0) * m (0, 1) - m (2, 1) * m (0, 0);
        n (0, 2) = m (0, 1) * m (1, 2) - m (0, 2) * m (1, 1);
        n (1, 2) = m (0, 2) * m (1, 0) - m (0, 0) * m (1, 2);
        n (2, 2) = m (0, 0) * m (1, 1) - m (0, 1) * m (1, 0);                          // 代数余子式

        Real_t det = m (0, 0) * n (0, 0) + m (0, 1) * n (1, 0) + m (0, 2) * n (2, 0);  // m行列式
        return det;
    }

    Real_t det4X4 (const Matrix& m)
    {
        // 先断定m的尺寸，否则报错
        if (m.GetRowSize () != 4 && m.GetColSize () != 4)
        {
            throw std::runtime_error ("Error: The matrix size is not 4*4");
        }
        Real_t   det = 0;
        Matrix submatrix (3, 3);  // 创建一个3x3矩阵，排除第0行和第i列

        submatrix (0, 0) = m (1, 1);
        submatrix (0, 1) = m (1, 2);
        submatrix (0, 2) = m (1, 3);
        submatrix (1, 0) = m (2, 1);
        submatrix (1, 1) = m (2, 2);
        submatrix (1, 2) = m (2, 3);
        submatrix (2, 0) = m (3, 1);
        submatrix (2, 1) = m (3, 2);
        submatrix (2, 2) = m (3, 3);

        det = m (0, 0) * det3X3 (submatrix);

        submatrix (0, 0) = m (1, 0);
        submatrix (0, 1) = m (1, 2);
        submatrix (0, 2) = m (1, 3);
        submatrix (1, 0) = m (2, 0);
        submatrix (1, 1) = m (2, 2);
        submatrix (1, 2) = m (2, 3);
        submatrix (2, 0) = m (3, 0);
        submatrix (2, 1) = m (3, 2);
        submatrix (2, 2) = m (3, 3);

        det -= m (0, 1) * det3X3 (submatrix);

        submatrix (0, 0) = m (1, 0);
        submatrix (0, 1) = m (1, 1);
        submatrix (0, 2) = m (1, 3);
        submatrix (1, 0) = m (2, 0);
        submatrix (1, 1) = m (2, 1);
        submatrix (1, 2) = m (2, 3);
        submatrix (2, 0) = m (3, 0);
        submatrix (2, 1) = m (3, 1);
        submatrix (2, 2) = m (3, 3);

        det += m (0, 2) * det3X3 (submatrix);

        submatrix (0, 0) = m (1, 0);
        submatrix (0, 1) = m (1, 1);
        submatrix (0, 2) = m (1, 2);
        submatrix (1, 0) = m (2, 0);
        submatrix (1, 1) = m (2, 1);
        submatrix (1, 2) = m (2, 2);
        submatrix (2, 0) = m (3, 0);
        submatrix (2, 1) = m (3, 1);
        submatrix (2, 2) = m (3, 2);

        det -= m (0, 3) * det3X3 (submatrix);

        return det;
    }

    Matrix identity (const std::size_t size)
    {
        Matrix I (size, size);
        for (std::size_t index = 0; index < size; ++index)
        {
            I (index, index) = 1.0;
        }

        return I;
    }

    /// @brief 生成向量的反对称矩阵
    Matrix skew (Real_t x, Real_t y, Real_t z)
    {
        Matrix W (3, 3);
        W (0, 1) = -z;
        W (0, 2) = y;
        W (1, 0) = z;
        W (1, 2) = -x;
        W (2, 0) = -y;
        W (2, 1) = x;

        return W;
    }

    Matrix hstack (const Matrix& m, const Matrix& n)
    {
        // 先判断m和n的行数是否相等
        if (m.GetRowSize () != n.GetRowSize ())
        {
            throw std::runtime_error ("Error: Matrices have different row sizes.");
        }

        Matrix result (m.GetRowSize (), m.GetColSize () + n.GetColSize ());
        result.BlockFill (0, 0, m);
        result.BlockFill (0, m.GetColSize (), n);
        return result;
    }

    Matrix vstack (const Matrix& m, const Matrix& n)
    {
        // 先判断m和n的列数是否相等
        if (m.GetColSize () != n.GetColSize ())
        {
            throw std::runtime_error ("Error: Matrices have different column sizes.");
        }

        Matrix result (m.GetRowSize () + n.GetRowSize (), m.GetColSize ());
        result.BlockFill (0, 0, m);
        result.BlockFill (m.GetRowSize (), 0, n);
        return result;
    }

    Matrix remove (const Matrix& m, const std::size_t Index, const std::size_t axis)
    {
        Matrix result;
        if (axis == 0)
        {  // 先判断Index是否大于矩阵的行数或者列数
            if (Index >= m.GetRowSize ())
            {
                throw std::runtime_error ("Error: Index is out of row range.");
            }

            result = Matrix (m.GetRowSize () - 1, m.GetColSize ());

            for (std::size_t i = 0; i < Index; ++i)
            {
                for (std::size_t j = 0; j < m.GetColSize (); ++j)
                {
                    result (i, j) = m (i, j);
                }
            }
            if (Index < m.GetRowSize () - 1)
            {
                for (std::size_t i = Index; i < result.GetRowSize (); ++i)
                {
                    for (std::size_t j = 0; j < m.GetColSize (); ++j)
                    {
                        result (i, j) = m (i + 1, j);
                    }
                }
            }
        }
        else if (axis == 1)
        {
            // 先判断Index是否大于行列式列数
            if (Index >= m.GetColSize ())
            {
                throw std::runtime_error ("Error: Index is out of column range.");
            }

            result = Matrix (m.GetRowSize (), m.GetColSize () - 1);

            for (std::size_t j = 0; j < Index; ++j)
            {
                for (std::size_t i = 0; i < m.GetRowSize (); ++i)
                {
                    result (i, j) = m (i, j);
                }
            }
            if (Index < m.GetColSize () - 1)
            {
                for (std::size_t j = Index; j < result.GetColSize (); ++j)
                {
                    for (std::size_t i = 0; i < m.GetRowSize (); ++i)
                    {
                        result (i, j) = m (i, j + 1);
                    }
                }
            }
        }
        return result;
    }

    /// @brief 利用 Eigen 库计算 3*3 实对称阵的特征值（升序）和特征向量
    void eigenValVec3X3 (const Matrix& m, _OUT Matrix& eigenValue, _OUT Matrix& eigenVector)
    {
        Eigen::Matrix3d eigenM;
        for (std::size_t i{0}; i < 3; ++i)
        {
            for (std::size_t j{0}; j < 3; ++j)
            {
                eigenM (i, j) = m (i, j);
            }
        }
        Eigen::SelfAdjointEigenSolver<Eigen::Matrix3d> eigenSolver (eigenM);
        if (eigenSolver.info () != Eigen::Success)
        {
            abort ();
        }
        Eigen::Vector3d eigenVal = eigenSolver.eigenvalues ().real ();  // 特征值默认按升序排列
        Eigen::Matrix3d eigenVec = eigenSolver.eigenvectors ().real ();
    
        for (std::size_t i{0}; i < 3; ++i)
        {
            eigenValue (i, 0) = eigenVal[i];
            for (std::size_t j{0}; j < 3; ++j)
            {
                eigenVector (i, j) = eigenVec (i, j);
            }
        }
    }

    Matrix orthonormalize (const Matrix& basis)
    {
        // 参考基矢量的 Gram 矩阵
        const auto gram = basis * basis.Transpose ();
    
        // Gram 矩阵特征值分解（特征值升序）
        Matrix eigenValue (3, 1);
        Matrix eigenVector (3, 3);
        eigenValVec3X3 (gram, eigenValue, eigenVector);
    
        // 旋转变换到特征（坐标系）空间，正交化
        auto orthBasis = eigenVector.Transpose () * basis;
    
        // 缩放变换（考虑最小特征值接近 0.0），标准正交化
        eigenValue (1, 0) = 1.0 / std::sqrt (eigenValue (1, 0));
        eigenValue (2, 0) = 1.0 / std::sqrt (eigenValue (2, 0));
        const auto det = det3X3 (orthBasis);
        Real_t lambda23 = eigenValue (1, 0) * eigenValue (2, 0);
        if (det < 0.0)
        {
            lambda23 = -lambda23;
        }
        orthBasis (0, 0) = (orthBasis (1, 1) * orthBasis (2, 2) - orthBasis (1, 2) * orthBasis (2, 1)) * lambda23;
        orthBasis (0, 1) = (orthBasis (1, 2) * orthBasis (2, 0) - orthBasis (1, 0) * orthBasis (2, 2)) * lambda23;
        orthBasis (0, 2) = (orthBasis (1, 0) * orthBasis (2, 1) - orthBasis (1, 1) * orthBasis (2, 0)) * lambda23;
        orthBasis (1, 0) *= eigenValue (1, 0);
        orthBasis (1, 1) *= eigenValue (1, 0);
        orthBasis (1, 2) *= eigenValue (1, 0);
        orthBasis (2, 0) *= eigenValue (2, 0);
        orthBasis (2, 1) *= eigenValue (2, 0);
        orthBasis (2, 2) *= eigenValue (2, 0);
    
        // 旋转变换逆变换，回到原来的坐标系空间
        return eigenVector * orthBasis;
    }

    bool operator== (const Matrix& a, const Matrix& b)
    {
        if (a.GetColSize () != b.GetColSize () || a.GetRowSize () != b.GetRowSize ())
        {
            std::cout << "matrix A's Col " << a.GetColSize () << " is not equal to " << "matrix B's Col " << b.GetColSize () << "\n";
            std::cout << "matrix A's Row " << a.GetRowSize () << " is not equal to " << "matrix B's Row " << b.GetRowSize () << "\n";
            return false;
        }

        const auto rowSize = a.GetRowSize ();
        const auto colSize = a.GetColSize ();
        for (std::size_t i = 0; i < rowSize; ++i)
        {
            for (std::size_t j = 0; j < colSize; ++j)
            {
                if (!(std::fabs (a (i, j) - b (i, j)) < TOLERANCE || std::fabs (a (i, j) - b (i, j)) < TOLERANCE * fabs (a (i, j))))
                {
                    std::cout << " when i & j = " << i << " " << j << ",  " << a (i, j) << " != " << b (i, j) << "\n";
                    return false;
                }
            }
        }

        return true;
    }
}  // namespace SG::Algebra

Matrix::Matrix (std::size_t rowSize, std::size_t colSize, Real_t value)
    : m_rowSize (rowSize), m_colSize (colSize), m_elements (rowSize * colSize, value)
{}

Matrix::Matrix (std::size_t rowSize, std::size_t colSize, const std::vector<Real_t>& elements, Real_t value)
{
    m_rowSize  = rowSize;
    m_colSize  = colSize;
    m_elements = elements;
    m_elements.resize (m_colSize * m_rowSize, value);
}

Matrix& Matrix::operator+= (const Matrix& b)
{
    assert (m_rowSize == b.m_rowSize && m_colSize == b.m_colSize);
    const auto size = m_rowSize * m_colSize;
    for (std::size_t index = 0; index < size; ++index)
    {
        m_elements[index] += b.m_elements[index];
    }

    return *this;
}

Matrix Matrix::operator+ (const Matrix& b) const
{
    assert (m_rowSize == b.m_rowSize && m_colSize == b.m_colSize);
    Matrix     result (m_rowSize, m_colSize);
    const auto size = m_rowSize * m_colSize;
    for (std::size_t index = 0; index < size; ++index)
    {
        result.m_elements[index] = m_elements[index] + b.m_elements[index];
    }

    return result;
}

Matrix& Matrix::operator-= (const Matrix& b)
{
    assert (m_rowSize == b.m_rowSize && m_colSize == b.m_colSize);
    const auto size = m_rowSize * m_colSize;
    for (std::size_t index = 0; index < size; ++index)
    {
        m_elements[index] -= b.m_elements[index];
    }

    return *this;
}

Matrix Matrix::operator- (const Matrix& b) const
{
    assert (m_rowSize == b.m_rowSize && m_colSize == b.m_colSize);
    Matrix     result (m_rowSize, m_colSize);
    const auto size = m_rowSize * m_colSize;
    for (std::size_t index = 0; index < size; ++index)
    {
        result.m_elements[index] = m_elements[index] - b.m_elements[index];
    }

    return result;
}

Matrix Matrix::operator* (const Real_t factor) const
{
    Matrix result{ *this };

    for (auto& ele : result.m_elements)
    {
        ele *= factor;
    }

    return result;
}

/** @brief 矩阵乘法
 *
 * 使用指针位移，避免通过 重载的括号运算符（i， j），此运算符内部使用乘法计算地址。
 * 指针偏移使用加法，效率高
 */
Matrix Matrix::operator* (const Matrix& b) const
{
    assert (this->m_colSize == b.m_rowSize);
    auto       row = this->m_rowSize;
    auto       col = b.m_colSize;
    Matrix     c (row, col);
    const auto m = this->m_colSize;

    Real_t* cPtr = &(c.m_elements[0]);
    for (std::size_t i = 0; i < row; ++i)
    {
        for (std::size_t j = 0; j < col; ++j)
        {
            // 将 bPtr 指向 b(0, j)
            const Real_t* bPtr = &b.m_elements[j];
            // 重新将 aPtr 指向其第 i 行第一个元素
            const Real_t* aPtr = &m_elements[i * this->m_colSize];

            Real_t sum = 0.0;
            for (std::size_t k = 0; k < m; ++k)
            {
                sum += (*bPtr) * (*aPtr);
                bPtr += col;  // 将 bPtr 下移一行
                ++aPtr;       // 将 aPtr 右移一列
            }
            // std::cout<< sum << "\n";
            *cPtr = sum;
            ++cPtr;  // 指向 c 的下一个要计算的元素
        }
    }

    return c;
}

Matrix Matrix::Transpose () const
{
    Matrix     result (m_colSize, m_rowSize);
    Real_t*    resPtr = &(result.m_elements[0]);
    const auto row    = m_colSize;
    const auto col    = m_rowSize;
    for (std::size_t i = 0; i < row; ++i)
    {
        // 将 bPtr 指向 b(0, i)
        const Real_t* bPtr = &(m_elements[i]);
        for (std::size_t j = 0; j < col; ++j)
        {
            // std::cout << "a " << *bPtr <<"\n";
            *resPtr = *bPtr;
            bPtr += m_colSize;
            ++resPtr;
        }
    }

    return result;
}

/// @brief 使用 LU 分解计算矩阵逆
Matrix Matrix::Inverse () const
{
    assert (m_colSize == m_rowSize);
    const std::size_t n = m_colSize;

    Matrix L (n, n);
    for (std::size_t i = 0; i < n; ++i)
    {
        L (i, i) = 1.0;
    }
    Matrix U = *this;

    Real_t factor{ 0.0 };
    for (std::size_t i = 0; i < n - 1; ++i)
    {
        factor = 1.0 / U (i, i);  // 假设主元不会出现零
        for (std::size_t j = i + 1; j < n; ++j)
        {
            if (fabs (U (j, i)) < 1e-12)
            {
                continue;
            }
            else
            {
                L (j, i) = U (j, i) * factor;
                for (std::size_t k = i; k < n; ++k)
                {
                    U (j, k) -= U (i, k) * L (j, i);
                }
            }
        }
    }

    // 计算 L 的逆矩阵
    Matrix Linv (n, n);
    for (std::size_t i = 0; i < n; ++i)
    {
        Linv (i, i) = 1.0;               // 考虑到 L(i, i) = 1.0
    }

    for (std::size_t r = 1; r < n; ++r)  // 列标
    {
        for (std::size_t c = 0; c < n - r; ++c)
        {
            // L(r+c, :) *Linv(:, r+c-1) = 0;
            Real_t sum = 0;
            for (std::size_t k = 0; k < r + c; k++)
            {
                sum += L (r + c, k) * Linv (k, c);
            }
            Linv (c + r, c) = -sum;
        }
    }

    // 计算上三角矩阵 U 的逆矩阵 （使用初等变换方法）
    Matrix Uinv (n, n);
    for (std::size_t i = 0; i < n; ++i)
    {
        Uinv (i, i) = 1.0;
    }
    // 初等变换，达到将 U 对角线变成 1.0
    for (std::size_t i = 0; i < n; ++i)
    {
        factor = 1.0 / U (i, i);
        for (std::size_t j = i; j < n; ++j)
        {
            U (i, j) *= factor;
            Uinv (i, j) *= factor;
        }
    }

    for (std::size_t i = 1; i < n; ++i)
    {
        for (std::size_t j = 0; j < i; ++j)
        {
            factor = -U (j, i);
            // 将第 i 行的 -U(j,i) 倍加到第 j 行, 使得 (j,i) 变成 0.0
            for (std::size_t k = i; k < n; ++k)
            {
                U (j, k) += factor * U (i, k);
                Uinv (j, k) += factor * Uinv (i, k);
            }
        }
    }

    return Uinv * Linv;
}

Matrix Matrix::Block (std::size_t i, std::size_t j, std::size_t m, std::size_t n)
{
    assert (this->m_rowSize > (i + m - 1));
    assert (this->m_colSize > (j + n - 1));

    Matrix block (m, n);
    for (std::size_t ii{0}; ii < m; ++ii)
    {
        for (std::size_t jj{0}; jj < n; ++jj)
        {
            block (ii, jj) = (*this) (i + ii, j + jj);
        }
    }

    return block;
}

Matrix& Matrix::BlockFill (const std::vector<std::size_t>& rows, const std::vector<std::size_t>& cols, const Matrix& b)
{
    std::size_t r = 0;
    for (const auto row : rows)
    {

        assert (row < m_rowSize);
        std::size_t n{0};
        for (const auto& col : cols)
        {
            (*this) (row, col) = b (r, n);
            ++n;
        }
        ++r;
    }

    return *this;
}

Matrix& Matrix::BlockFill (std::size_t rowStart, std::size_t colStart, const Matrix& b)
{
    assert (this->m_colSize > (colStart + b.GetColSize () - 1));
    assert (this->m_rowSize > (rowStart + b.GetRowSize () - 1));

    const auto nbRow = b.GetRowSize ();
    const auto nbCol = b.GetColSize ();
    for (std::size_t r = 0; r < nbRow; ++r)
    {
        const auto row = rowStart + r;
        for (std::size_t n = 0; n < nbCol; ++n)
        {
            const auto col     = colStart + n;
            (*this) (row, col) = b (r, n);
        }
    }

    return *this;
}

Matrix& Matrix::BlockSum (const std::vector<std::size_t>& rows, const std::vector<std::size_t>& cols, const Matrix& sub)
{
    std::size_t r = 0;
    for (const auto row : rows)
    {

        assert (row < m_rowSize);
        std::size_t n{0};
        for (const auto& col : cols)
        {
            (*this) (row, col) += sub (r, n);
            ++n;
        }
        ++r;
    }

    return *this;
}